/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.lucene.search;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import org.apache.lucene.util.FloatComparator;
import org.apache.lucene.util.PriorityQueue;

/** Base class for Scorers that score disjunctions. */
abstract class DisjunctionScorer extends Scorer {

  private final int numClauses;
  private final boolean needsScores;

  private final DisjunctionDISIApproximation approximation;
  private final TwoPhase twoPhase;

  protected DisjunctionScorer(List<Scorer> subScorers, ScoreMode scoreMode, long leadCost)
      throws IOException {
    if (subScorers.size() <= 1) {
      throw new IllegalArgumentException("There must be at least 2 subScorers");
    }
    this.numClauses = subScorers.size();
    this.needsScores = scoreMode != ScoreMode.COMPLETE_NO_SCORES;
    boolean hasApproximation = false;
    float sumMatchCost = 0;
    long sumApproxCost = 0;
    List<DisiWrapper> wrappers = new ArrayList<>();
    for (Scorer scorer : subScorers) {
      DisiWrapper w = new DisiWrapper(scorer, false);
      long costWeight = (w.cost <= 1) ? 1 : w.cost;
      sumApproxCost += costWeight;
      if (w.twoPhaseView != null) {
        hasApproximation = true;
        sumMatchCost += w.matchCost * costWeight;
      }
      wrappers.add(w);
    }
    this.approximation = new DisjunctionDISIApproximation(wrappers, leadCost);

    if (hasApproximation == false) { // no sub scorer supports approximations
      twoPhase = null;
    } else {
      final float matchCost = sumMatchCost / sumApproxCost;
      twoPhase = new TwoPhase(approximation, matchCost);
    }
  }

  @Override
  public DocIdSetIterator iterator() {
    if (twoPhase != null) {
      return TwoPhaseIterator.asDocIdSetIterator(twoPhase);
    } else {
      return approximation;
    }
  }

  @Override
  public TwoPhaseIterator twoPhaseIterator() {
    return twoPhase;
  }

  private class TwoPhase extends TwoPhaseIterator {

    private final float matchCost;
    // list of verified matches on the current doc
    DisiWrapper verifiedMatches;
    // priority queue of approximations on the current doc that have not been verified yet
    final PriorityQueue<DisiWrapper> unverifiedMatches;

    private TwoPhase(DocIdSetIterator approximation, float matchCost) {
      super(approximation);
      this.matchCost = matchCost;
      unverifiedMatches =
          PriorityQueue.usingComparator(numClauses, FloatComparator.comparing(d -> d.matchCost));
    }

    DisiWrapper getSubMatches() throws IOException {
      // iteration order does not matter
      for (DisiWrapper w : unverifiedMatches) {
        if (w.twoPhaseView.matches()) {
          w.next = verifiedMatches;
          verifiedMatches = w;
        }
      }
      unverifiedMatches.clear();
      assert assertVerifiedMatchesApproximationDocID();
      return verifiedMatches;
    }

    private boolean assertVerifiedMatchesApproximationDocID() {
      int docID = docID();
      for (DisiWrapper w = verifiedMatches; w != null; w = w.next) {
        if (w.doc != docID) {
          return false;
        }
      }
      return true;
    }

    @Override
    public boolean matches() throws IOException {
      verifiedMatches = null;
      unverifiedMatches.clear();

      for (DisiWrapper w = DisjunctionScorer.this.approximation.topList(); w != null; ) {
        DisiWrapper next = w.next;

        if (w.twoPhaseView == null) {
          // implicitly verified, move it to verifiedMatches
          w.next = verifiedMatches;
          verifiedMatches = w;

          if (needsScores == false) {
            // we can stop here
            return true;
          }
        } else {
          unverifiedMatches.add(w);
        }
        w = next;
      }

      if (verifiedMatches != null) {
        return true;
      }

      // verify subs that have an two-phase iterator
      // least-costly ones first
      while (unverifiedMatches.size() > 0) {
        DisiWrapper w = unverifiedMatches.pop();
        if (w.twoPhaseView.matches()) {
          w.next = null;
          verifiedMatches = w;
          return true;
        }
      }

      return false;
    }

    @Override
    public float matchCost() {
      return matchCost;
    }
  }

  @Override
  public final int docID() {
    return approximation.docID();
  }

  DisiWrapper getSubMatches() throws IOException {
    if (twoPhase == null) {
      return approximation.topList();
    } else {
      return twoPhase.getSubMatches();
    }
  }

  @Override
  public final float score() throws IOException {
    return score(getSubMatches());
  }

  /** Compute the score for the given linked list of scorers. */
  protected abstract float score(DisiWrapper topList) throws IOException;

  @Override
  public final Collection<ChildScorable> getChildren() throws IOException {
    ArrayList<ChildScorable> children = new ArrayList<>();
    for (DisiWrapper scorer = getSubMatches(); scorer != null; scorer = scorer.next) {
      children.add(new ChildScorable(scorer.scorer, "SHOULD"));
    }
    return children;
  }
}
